"""
Ollama API客户端模块 - 使用Python标准库
"""

import base64
import json
import urllib.request
import urllib.error
from pathlib import Path
from typing import Optional, Dict, Any
import time


class OllamaClient:
    """Ollama API客户端"""
    
    def __init__(self, host: str = "http://localhost:11434", model: str = "qwen2.5-vl:7b", system_prompt: str = None):
        """
        初始化Ollama客户端
        
        Args:
            host: Ollama服务地址
            model: 使用的模型名称
            system_prompt: 系统提示词，如果为None则使用默认提示词
        """
        self.host = host.rstrip('/')
        self.model = model
        self.api_url = f"{self.host}/api/generate"
        
        # 系统提示词
        if system_prompt:
            self.system_prompt = system_prompt
        else:
            # 默认系统提示词
            self.system_prompt = """你是一个图像内容分析工具。请基于图片，仅对其中可见的、客观的元素进行事实性描述，并最终整合成一段连贯、中性的文字段落。分析焦点（请根据以下要点组织你的观察，但最终输出必须是段落形式）：
人物（核心）：
外貌特征： 性别、大致年龄、发型与发色、面部朝向。
衣着： 服装的款式、颜色、类型。
姿态与动作： 站姿、坐姿或具体动作，手臂和手的位置。
背景环境：
场景： 例如海滩、房间、街道等。
物体与元素： 存在的具体物体，如沙子、海浪、云、树木、家具等。
颜色与光照： 主导颜色（如蓝天、黄沙）、光照条件（如白天、阴影）。
严格禁止事项：
禁止添加任何主观感受、情感、氛围或评价。 例如，绝对不允许出现"宁静自然"、"自由放松"、"美丽"、"表情平静"等词语。
禁止描述图片中不可见或无法直接推断的内容。
禁止使用比喻、象征或任何文学性修辞手法。
禁止使用项目符号、数字编号或分点列表。
输出必须是一个完整的段落。
输出要求： 请直接以"图片中..."或"图中..."开头，生成一段纯粹客观的描述性文字。"""
    
    def encode_image_to_base64(self, image_path: Path) -> str:
        """
        将图片文件编码为base64字符串
        
        Args:
            image_path: 图片文件路径
            
        Returns:
            base64编码的图片字符串
        """
        try:
            with open(image_path, "rb") as image_file:
                return base64.b64encode(image_file.read()).decode('utf-8')
        except Exception as e:
            raise Exception(f"读取图片文件失败: {image_path} - {str(e)}")
    
    def check_connection(self) -> bool:
        """
        检查与Ollama服务的连接
        
        Returns:
            连接是否正常
        """
        try:
            req = urllib.request.Request(f"{self.host}/api/tags")
            with urllib.request.urlopen(req, timeout=10) as response:
                return response.status == 200
        except Exception:
            return False
    
    def check_model_available(self) -> bool:
        """
        检查指定模型是否可用
        
        Returns:
            模型是否可用
        """
        try:
            req = urllib.request.Request(f"{self.host}/api/tags")
            with urllib.request.urlopen(req, timeout=10) as response:
                if response.status == 200:
                    data = json.loads(response.read().decode('utf-8'))
                    models = data.get('models', [])
                    return any(model.get('name', '').startswith(self.model.split(':')[0]) for model in models)
            return False
        except Exception:
            return False
    
    def analyze_image(self, image_path: Path, max_retries: int = 3) -> Optional[str]:
        """
        使用Qwen2.5VL模型分析图片内容
        
        Args:
            image_path: 图片文件路径
            max_retries: 最大重试次数
            
        Returns:
            图片描述文本，失败时返回None
        """
        for attempt in range(max_retries):
            try:
                # 编码图片
                image_base64 = self.encode_image_to_base64(image_path)
                
                # 构建请求数据
                payload = {
                    "model": self.model,
                    "prompt": self.system_prompt,
                    "images": [image_base64],
                    "stream": False,
                    "options": {
                        "temperature": 0.1,  # 降低温度以获得更客观的描述
                        "max_tokens": 1000
                    }
                }
                
                print(f"正在分析图片: {image_path.name} (尝试 {attempt + 1}/{max_retries})")
                
                # 发送请求
                data = json.dumps(payload).encode('utf-8')
                req = urllib.request.Request(
                    self.api_url,
                    data=data,
                    headers={'Content-Type': 'application/json'}
                )
                
                with urllib.request.urlopen(req, timeout=120) as response:
                    if response.status == 200:
                        result = json.loads(response.read().decode('utf-8'))
                        description = result.get('response', '').strip()
                        
                        if description:
                            print(f"分析完成: {image_path.name}")
                            return description
                        else:
                            print(f"警告: 模型返回空描述: {image_path.name}")
                            return None
                    else:
                        print(f"API请求失败 (状态码: {response.status}): {image_path.name}")
                        if attempt < max_retries - 1:
                            print(f"等待 {2 ** attempt} 秒后重试...")
                            time.sleep(2 ** attempt)
                        
            except urllib.error.URLError as e:
                print(f"网络错误: {image_path.name} - {str(e)}")
                if attempt < max_retries - 1:
                    print(f"等待 {2 ** attempt} 秒后重试...")
                    time.sleep(2 ** attempt)
            except Exception as e:
                print(f"分析图片时发生错误: {image_path.name} - {str(e)}")
                if attempt < max_retries - 1:
                    print(f"等待 {2 ** attempt} 秒后重试...")
                    time.sleep(2 ** attempt)
        
        print(f"分析失败，已达到最大重试次数: {image_path.name}")
        return None
    
    def get_model_info(self) -> Optional[Dict[str, Any]]:
        """
        获取模型信息
        
        Returns:
            模型信息字典
        """
        try:
            req = urllib.request.Request(f"{self.host}/api/tags")
            with urllib.request.urlopen(req, timeout=10) as response:
                if response.status == 200:
                    data = json.loads(response.read().decode('utf-8'))
                    models = data.get('models', [])
                    for model in models:
                        if model.get('name', '').startswith(self.model.split(':')[0]):
                            return model
            return None
        except Exception as e:
            print(f"获取模型信息失败: {str(e)}")
            return None